from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
import random

# 1. 连接到 Milvus 服务
def connect_to_milvus():
    try:
        connections.connect(
            alias="default",
            host='127.0.0.1',
            port='19530'
        )
        print("Connected to Milvus successfully!")
    except Exception as e:
        print(f"Failed to connect to Milvus: {e}")

# 2. 创建集合
def create_collection():
    # 定义字段
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=128)
    ]
    # 定义集合模式
    schema = CollectionSchema(fields=fields, description="My first Milvus collection")
    # 创建集合
    collection = Collection(name="my_collection", schema=schema)
    print("Collection created successfully!")
    return collection

# 3. 插入向量数据
def insert_data(collection):
    # 生成随机向量数据
    vectors = [[random.random() for _ in range(128)] for _ in range(100)]
    # 插入数据
    data = [vectors]
    insert_result = collection.insert(data)
    collection.flush()
    print(f"Inserted {insert_result.insert_count} rows of data.")

# 4. 创建索引
def create_index(collection):
    index = {
        "index_type": "IVF_FLAT",
        "metric_type": "L2",
        "params": {"nlist": 128}
    }
    collection.create_index(field_name="embedding", index_params=index)
    print("Index created successfully!")

# 5. 进行向量搜索
def search_vectors(collection):
    # 生成一个查询向量
    query_vector = [random.random() for _ in range(128)]
    search_params = {
        "metric_type": "L2",
        "params": {"nprobe": 10}
    }
    results = collection.search(
        data=[query_vector],
        anns_field="embedding",
        param=search_params,
        limit=10,
        output_fields=["id"]
    )
    for hit in results[0]:
        print(f"ID: {hit.id}, Distance: {hit.distance}")

# 主函数
def main():
    # 连接到 Milvus
    connect_to_milvus()
    # 创建集合
    collection = create_collection()
    # 插入数据
    insert_data(collection)
    collection.load();
    # 创建索引
    create_index(collection)
    # 进行向量搜索
    search_vectors(collection)

main()